# coding=utf-8
import tensorflow as tf
import numpy as np
tf.enable_eager_execution()
n_samples = 500
learning_rate = 0.01
training_steps = 5000
display_step = 1000

#Create dataset
X = np.random.rand(n_samples).astype(np.float32)
Y = X * 10 + 5
W = tf.Variable(tf.random.normal([1]))
b = tf.Variable(tf.zeros([1]))


# Define LR and Loss function (MSE)
def linear_regression(x):
    return W * x + b


def mean_square(y_pred, y_true):
    return tf.reduce_sum(tf.pow(y_pred - y_true, 2)) / (2 * n_samples)


# Stochastic Gradient Descent Optimizer.
optimizer = tf.compat.v2.optimizers.SGD(learning_rate)


def run_optimization(step):
    # Wrap computation inside a GradientTape for automatic differentiation.
    with tf.GradientTape() as g:
        pred = linear_regression(X)
        loss = mean_square(pred, Y)

    # Compute gradients.
    gradients = g.gradient(loss, [W, b])

    # Update W and b following gradients.
    optimizer.apply_gradients(zip(gradients, [W, b]))
    if step % display_step == 0:
        print("step: %i, loss: %f, W: %f, b: %f" % (step, loss, W.numpy(), b.numpy()))

for step in range(1, training_steps + 1):
    # Run the optimization to update W and b values.
    run_optimization(step)
    # if step % display_step == 0:
    #     pred = linear_regression(X)
    #     loss = mean_square(pred, Y)
    #     print("step: %i, loss: %f, W: %f, b: %f" % (step, loss, W.numpy(), b.numpy()))